package org.example.doris.arrow.jdbc;

import org.apache.arrow.adbc.core.AdbcConnection;
import org.apache.arrow.adbc.core.AdbcDatabase;
import org.apache.arrow.adbc.core.AdbcDriver;
import org.apache.arrow.adbc.core.AdbcStatement;
import org.apache.arrow.adbc.driver.jdbc.JdbcDriver;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowReader;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

public class JDBCDriverDemo {
    public static void main(String[] args) throws Exception {
        //Class.forName("org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver");
        String DB_URL = "jdbc:arrow-flight-sql://192.168.0.227:9090?useServerPrepStmts=false"
                + "&cachePrepStmts=true&useSSL=false&useEncryption=false";
        final HashMap<String, Object> parameters = new HashMap<>();
        AdbcDriver.PARAM_URI.set(parameters, DB_URL);
        AdbcDriver.PARAM_USERNAME.set(parameters, "root");
        AdbcDriver.PARAM_PASSWORD.set(parameters, "123456");


        BufferAllocator allocator = new RootAllocator();
        AdbcDatabase db = new JdbcDriver(allocator).open(parameters);
        AdbcConnection connection = db.connect();
        AdbcStatement stmt = connection.createStatement();

        stmt.setSqlQuery("select * from testdb.table_hash;");
        AdbcStatement.QueryResult queryResult = stmt.executeQuery();
        ArrowReader reader = queryResult.getReader();
        List<String> result = new ArrayList<>();
        while (reader.loadNextBatch()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            String tsvString = root.contentToTSVString();
            result.add(tsvString);
        }
        long etime = System.currentTimeMillis();
        System.out.printf("batchs %d\n", result.size());

        reader.close();
        queryResult.close();
        stmt.close();

    }
}
